{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "### Exploratory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import rp\n",
    "import numpy as np\n",
    "from icecream import ic \n",
    "from einops import rearrange, repeat\n",
    "import torch\n",
    "\n",
    "def get_video_mask(video):\n",
    "    \"\"\" Given a THWC video, return a THW binary mask \"\"\"\n",
    "    assert rp.is_numpy_array(video)\n",
    "    num_frames,height,width,num_channels=video.shape\n",
    "\n",
    "    mask = np.zeros((num_frames,height,width),dtype=bool)\n",
    "    \n",
    "    #Binary matrix with shape (height, width)\n",
    "    circle = rp.flat_circle_kernel(diameter=min(height,width))\n",
    "    circle = rp.crop_image(circle,height,width,origin='center')\n",
    "    circle = rp.as_binary_image(circle)\n",
    "    mask[0]=circle\n",
    "    # mask[5]=circle\n",
    "    mask[len(mask)-1]=circle\n",
    "\n",
    "    return mask\n",
    "\n",
    "def get_all_mask_positions(mask):\n",
    "    \"\"\"\n",
    "    Get (t, x, y) coordinates of all True pixels in the mask.\n",
    "    \"\"\"\n",
    "    assert mask.ndim==3\n",
    "    assert mask.dtype==bool\n",
    "    \n",
    "    ts, ys, xs= np.where(mask)\n",
    "    #NOTE: np.where returns ys then xs, we return xs then ys!\n",
    "    \n",
    "    return np.stack([ts, xs, ys]).T # M×3\n",
    "\n",
    "def get_random_mask_positions(mask, N):\n",
    "    \"\"\"\n",
    "    Returns size [N, TXY]\n",
    "    \"\"\"\n",
    "    assert mask.ndim==3\n",
    "    assert mask.dtype==bool\n",
    "    T,H,W = mask.shape\n",
    "    \n",
    "    all_positions = get_all_mask_positions(mask) # M×3\n",
    "    chosen_positions = rp.random_batch_with_replacement(all_positions, N) #N×3\n",
    "    chosen_positions = rp.as_numpy_array(chosen_positions)\n",
    "\n",
    "    assert chosen_positions.shape == (N, 3)\n",
    "    return chosen_positions #N×3\n",
    "\n",
    "def get_random_mask_positions_per_frame(mask, N):\n",
    "    T,H,W = mask.shape\n",
    "    \n",
    "    chosen_positions = get_random_mask_positions(mask, N)\n",
    "\n",
    "    output = [[] for _ in range(T)]\n",
    "    for t,x,y in chosen_positions:\n",
    "        output[t].append([x,y])\n",
    "    \n",
    "    return output\n",
    "\n",
    "def demo_get_random_mask_positions_per_frame():\n",
    "    video = rp.load_video(\n",
    "        \"https://github.com/facebookresearch/co-tracker/raw/refs/heads/main/assets/apple.mp4\",\n",
    "        use_cache=True,\n",
    "    )\n",
    "    video = rp.resize_images(video, size=(480,720))\n",
    "    video = rp.as_numpy_array(video)\n",
    "    \n",
    "    mask = get_video_mask(video)\n",
    "    \n",
    "    mask_positions_per_frame = get_random_mask_positions_per_frame(mask, 1000)\n",
    "    \n",
    "    dotted_video = []\n",
    "    for mask_positions, frame in zip(mask_positions_per_frame, video):\n",
    "        if len(mask_positions):\n",
    "            x, y = rp.list_transpose(mask_positions) #Nx3 -> 3xN\n",
    "            frame = rp.cv_draw_circles(frame, x, y, color='light blue', rim=2)\n",
    "        dotted_video.append(frame)\n",
    "    \n",
    "    rp.display_video(rp.horizontally_concatenated_videos(mask,dotted_video), framerate=10)\n",
    "\n",
    "demo_get_random_mask_positions_per_frame()\n",
    "\n",
    "device = rp.select_torch_device(prefer_used=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "video = rp.load_video(\n",
    "    \"https://github.com/facebookresearch/co-tracker/raw/refs/heads/main/assets/apple.mp4\",\n",
    "    use_cache=True,\n",
    ")\n",
    "video = rp.resize_images(video, size=(480,720))\n",
    "video = rp.as_numpy_array(video)\n",
    "\n",
    "\n",
    "import dinkydata\n",
    "\n",
    "\n",
    "mask = get_video_mask(video)\n",
    "\n",
    "mask_positions_per_frame = get_random_mask_positions_per_frame(mask, 1000)\n",
    "\n",
    "dotted_video = []\n",
    "for mask_positions, frame in zip(mask_positions_per_frame, video):\n",
    "    if len(mask_positions):\n",
    "        x, y = rp.list_transpose(mask_positions) #Nx3 -> 3xN\n",
    "        frame = rp.cv_draw_circles(frame, x, y, color='light blue', rim=2)\n",
    "    dotted_video.append(frame)\n",
    "\n",
    "rp.display_video(rp.horizontally_concatenated_videos(mask,dotted_video), framerate=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample 100 random points from the circle mask\n",
    "track_points = get_random_mask_positions(mask, 1000) #100 x 3\n",
    "\n",
    "\n",
    "# # Run cotracker with our custom points\n",
    "tracks, visibility = rp.run_cotracker(\n",
    "    video, \n",
    "    device=device,\n",
    "    queries=track_points,  # Use our custom initialization points\n",
    "    # grid_size = None,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "tracks.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "dotted_video = []\n",
    "for frame, track in zip(video, tracks):\n",
    "    x, y = track.T\n",
    "    frame = rp.cv_draw_circles(frame, x, y, color='light blue', rim=2, )\n",
    "    dotted_video.append(frame)\n",
    "rp.display_video(dotted_video)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Visualization functions\n",
    "# def generate_uv_colors(tracks, frame_idx=0):\n",
    "#     \"\"\"Generate colors based on UV coordinates of initial positions.\"\"\"\n",
    "#     height, width = video[0].shape[:2]\n",
    "    \n",
    "#     colors = []\n",
    "#     for x, y in tracks[frame_idx]:\n",
    "#         x_py = to_python_num(x)\n",
    "#         y_py = to_python_num(y)\n",
    "#         # Generate colors based on UV coordinates\n",
    "#         r = x_py/width  \n",
    "#         g = y_py/height\n",
    "#         b = 0\n",
    "#         colors.append((r, g, b))\n",
    "    \n",
    "#     return colors\n",
    "\n",
    "# def draw_tracks_on_video(video, tracks, colors, radius=5):\n",
    "#     \"\"\"Draw colored track points on video frames.\"\"\"\n",
    "#     result = []\n",
    "#     for frame_number, frame in enumerate(video):\n",
    "#         frame_copy = frame.copy()\n",
    "#         for color, (x, y) in zip(colors, tracks[frame_number]):\n",
    "#             x_py = to_python_num(x)\n",
    "#             y_py = to_python_num(y)\n",
    "            \n",
    "#             frame_copy = cv_draw_circle(\n",
    "#                 frame_copy,\n",
    "#                 x_py, y_py,\n",
    "#                 radius=radius,\n",
    "#                 color=color,\n",
    "#                 antialias=True,\n",
    "#                 copy=False,\n",
    "#             )\n",
    "#         result.append(frame_copy)\n",
    "#     return result\n",
    "\n",
    "# # Apply visualization\n",
    "# colors = generate_uv_colors(tracks)\n",
    "# new_video = draw_tracks_on_video(video, tracks, colors, radius=8)\n",
    "# display_video(new_video)\n",
    "# print(f\"SAVED: {save_video_mp4(new_video, 'circle_tracking.mp4')}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "track_points.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dinkydata.counterfactual_dataset.counterfactual_dataset as dd\n",
    "dd.demo()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rp import *\n",
    "table = dd.index_table()\n",
    "cartridge = rp.random_element(table).load()\n",
    "rp.display_video(cartridge.preview())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "video = cartridge.orig_vid\n",
    "gen_video = cartridge.gen_vid\n",
    "\n",
    "mask = rp.resize_list(cartridge.mask, 49)\n",
    "mask = rp.resize_images(mask, size=(480, 720))\n",
    "mask = rp.as_grayscale_images(rp.as_binary_images(mask))\n",
    "\n",
    "N=200\n",
    "mask_positions_per_frame = get_random_mask_positions_per_frame(mask, N)\n",
    "colors = [rp.as_rgb_float_color('random') for _ in range(N)]\n",
    "\n",
    "# dotted_video = []\n",
    "# for mask_positions, frame in zip(mask_positions_per_frame, video):\n",
    "#     if len(mask_positions):\n",
    "#         x, y = rp.list_transpose(mask_positions) #Nx3 -> 3xN\n",
    "#         frame = rp.cv_draw_circles(frame, x, y, color='light blue', rim=2)\n",
    "#     dotted_video.append(frame)\n",
    "\n",
    "# rp.display_video(rp.horizontally_concatenated_videos(mask,dotted_video), framerate=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample 100 random points from the circle mask\n",
    "track_points = get_random_mask_positions(mask, N) #100 x 3\n",
    "\n",
    "\n",
    "# # Run cotracker with our custom points\n",
    "tracks, visibility = rp.run_cotracker(\n",
    "    video, \n",
    "    device=\"cuda\",\n",
    "    queries=track_points,  # Use our custom initialization points\n",
    "    # grid_size = None,\n",
    ")\n",
    "\n",
    "\n",
    "# # Run cotracker with our custom points\n",
    "tracks2, visibility2 = rp.run_cotracker(\n",
    "    gen_video, \n",
    "    device=\"cuda\",\n",
    "    queries=track_points,  # Use our custom initialization points\n",
    "    # grid_size = None,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "dotted_video = []\n",
    "for frame, gen_frame, track, gen_track, vis, gen_vis, m in rp.eta(zip(\n",
    "    video, gen_video, tracks, tracks2, visibility, visibility2, mask\n",
    "),length=len(video)):\n",
    "    x, y = rp.as_numpy_array(track.T)\n",
    "    gx, gy = rp.as_numpy_array(gen_track.T)\n",
    "\n",
    "    \n",
    "    v = vis.T * gen_vis.T  # Only show when both are visibile\n",
    "    radii = v * 5 + 1\n",
    "    frame = rp.cv_draw_circles(frame, x, y, color=colors, rim=2, radius=radii)\n",
    "    gen_frame = rp.cv_draw_circles(gen_frame, gx, gy, color=colors, rim=2, radius=radii)\n",
    "    \n",
    "    avg_frame = frame // 2 + gen_frame // 2\n",
    "    avg_frame = rp.cv_draw_contours(avg_frame, [[[xi,yi],[xj,yj]] for xi,yi,xj,yj in zip(x,y,gx,gy)])\n",
    "    dotted_frame = rp.horizontally_concatenated_images(\n",
    "        rp.labeled_images(\n",
    "            [\n",
    "                frame,\n",
    "                gen_frame,\n",
    "                rp.blend_images(rp.blend_images(avg_frame, rp.auto_canny(m), mode=\"add\"),m,.1,mode='add'),\n",
    "            ],\n",
    "            [\n",
    "                \"Original Video\",\n",
    "                \"Generated Video\",\n",
    "                \"Blended (for visualiztion)\",\n",
    "            ],\n",
    "            font=\"R:Futura\",\n",
    "        )\n",
    "    )\n",
    "    dotted_video.append(dotted_frame)\n",
    "rp.display_video(dotted_video,framerate=10)\n",
    "rp.save_video_mp4(dotted_video, framerate=10)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "### Final Func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "import rp\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def fast_scatter_add(output_tensor, latent_tracks, track_colors, num_timesteps, num_points, width, height):\n",
    "    \"\"\"\n",
    "    Efficiently adds tracking point colors to a latent tensor using scatter_add.\n",
    "    \n",
    "    Analogy: \n",
    "    The original implementation is like manually placing colored dots on a canvas one at a time:\n",
    "    ```\n",
    "    for lt in range(LT):\n",
    "        for n in range(N):\n",
    "            color = track_colors[n]\n",
    "            x, y = latent_tracks[lt, n].long()\n",
    "            if 0<=x<LW and 0<=y<LH:\n",
    "                dotted_latent[lt, :, y, x] += color\n",
    "    ```\n",
    "    \n",
    "    This optimized version is like having a machine that can place all dots of the same color\n",
    "    simultaneously at their respective coordinates.\n",
    "    \n",
    "    Args:\n",
    "        output_tensor: Zero-initialized tensor to populate with values (LT, LC, LH, LW)\n",
    "        latent_tracks: Tensor containing track coordinates (LT, N, 2)\n",
    "        track_colors: Tensor of colors for each track point (N, LC)\n",
    "        num_timesteps: Number of timesteps (LT)\n",
    "        num_points: Number of track points (N)\n",
    "        width: Width of latent tensor (LW)\n",
    "        height: Height of latent tensor (LH)\n",
    "        \n",
    "    Returns:\n",
    "        torch.Tensor: Populated tensor with added track colors\n",
    "    \"\"\"\n",
    "    device = output_tensor.device\n",
    "    dtype = output_tensor.dtype\n",
    "    C = track_colors.shape[1]  # Number of channels\n",
    "    \n",
    "    # Extract x and y coordinates\n",
    "    xs, ys = latent_tracks[..., 0], latent_tracks[..., 1]\n",
    "    \n",
    "    # Create mask for valid coordinates (within bounds)\n",
    "    valid_mask = (xs >= 0) & (xs < width) & (ys >= 0) & (ys < height)\n",
    "    \n",
    "    # Loop over timesteps (still needed, but with vectorized inner operations)\n",
    "    for t in range(num_timesteps):\n",
    "        # Get valid points for this timestep\n",
    "        t_valid = valid_mask[t]\n",
    "        if not t_valid.any():\n",
    "            continue\n",
    "            \n",
    "        # Get coordinates of valid points\n",
    "        valid_xs = xs[t, t_valid]\n",
    "        valid_ys = ys[t, t_valid]\n",
    "        valid_indices = t_valid.nonzero().squeeze(1)\n",
    "        \n",
    "        # Get colors for valid points\n",
    "        valid_colors = track_colors[valid_indices]\n",
    "        \n",
    "        # Compute indices for scatter operation\n",
    "        indices = valid_ys * width + valid_xs\n",
    "        \n",
    "        # Perform scatter_add for each channel\n",
    "        for c in range(C):\n",
    "            # Create flat output tensor for this timestep and channel\n",
    "            flat_output = output_tensor[t, c].view(-1)\n",
    "            \n",
    "            # Scatter add the values\n",
    "            flat_output.scatter_add_(0, indices, valid_colors[:, c])\n",
    "    \n",
    "    return output_tensor\n",
    "\n",
    "def get_all_mask_positions(mask):\n",
    "    \"\"\"\n",
    "    Get (t, x, y) coordinates of all True pixels in the mask.\n",
    "    \n",
    "    Args:\n",
    "        mask: A 3D boolean mask of shape (T, H, W) where True indicates positions to extract.\n",
    "              Can be either torch.Tensor or numpy.ndarray.\n",
    "    \n",
    "    Returns:\n",
    "        torch.Tensor: Tensor of shape (M, 3) containing (t, x, y) coordinates of all True pixels,\n",
    "                      where M is the number of True pixels in the mask\n",
    "    \"\"\"\n",
    "    assert mask.ndim == 3\n",
    "    assert mask.dtype == bool or mask.dtype == torch.bool\n",
    "    \n",
    "    if isinstance(mask, torch.Tensor):\n",
    "        # Already a torch tensor, use torch.where\n",
    "        indices = torch.where(mask)\n",
    "        ts, ys, xs = indices\n",
    "        # Stack to get M×3 tensor\n",
    "        return torch.stack([ts, xs, ys], dim=1)\n",
    "    else:\n",
    "        # Convert numpy array to torch tensor\n",
    "        mask_tensor = torch.from_numpy(mask)\n",
    "        indices = torch.where(mask_tensor)\n",
    "        ts, ys, xs = indices\n",
    "        #NOTE: torch.where (like np.where) returns ys then xs, we return xs then ys!\n",
    "        return torch.stack([ts, xs, ys], dim=1)\n",
    "\n",
    "def get_random_mask_positions(mask, N):\n",
    "    \"\"\"\n",
    "    Randomly samples N positions from the True pixels in a 3D mask.\n",
    "    \n",
    "    Args:\n",
    "        mask: A 3D boolean mask of shape (T, H, W) where True indicates valid positions.\n",
    "              Can be either torch.Tensor or numpy.ndarray.\n",
    "        N: Number of random positions to sample with replacement\n",
    "    \n",
    "    Returns:\n",
    "        torch.Tensor: Tensor of shape (N, 3) containing randomly sampled (t, x, y) coordinates\n",
    "    \"\"\"\n",
    "    assert mask.ndim == 3\n",
    "    assert mask.dtype == bool or mask.dtype == torch.bool\n",
    "    T, H, W = mask.shape\n",
    "    \n",
    "    # Get all positions as a torch tensor\n",
    "    all_positions = get_all_mask_positions(mask)  # M×3 torch tensor\n",
    "    \n",
    "    if len(all_positions) == 0:\n",
    "        raise ValueError(\"Mask contains no True values to sample from\")\n",
    "    \n",
    "    # Determine device based on the tensor\n",
    "    device = all_positions.device\n",
    "    \n",
    "    # Randomly sample N indices with replacement\n",
    "    idx = torch.randint(0, len(all_positions), (N,), device=device)\n",
    "    chosen_positions = all_positions[idx]\n",
    "    \n",
    "    assert chosen_positions.shape == (N, 3)\n",
    "    return chosen_positions  # N×3 torch tensor\n",
    "\n",
    "def generate_dotted_latents(\n",
    "    videos,\n",
    "    latent_mask,\n",
    "    *,\n",
    "    out_channels = 16, \n",
    "    num_points=1024,\n",
    "    device = None,\n",
    "    silent = True,\n",
    "):\n",
    "    \"\"\"\n",
    "    Generates a latent-tracking-point control video from a source video\n",
    "    Uses CoTracker to move random dots around a latent video\n",
    "\n",
    "    Args:\n",
    "        - videos: Give a list of videos, either in torch BTCHW form or a list of rp videos\n",
    "        - latent_mask: Determines the duration and dimensions of the output, \n",
    "                       and where tracking points are initialized.\n",
    "                       Either numpy arrays or torch tensors are ok.\n",
    "                       In THW form. \n",
    "        - out_channels: Determines the number of channels in the output\n",
    "        - num_points: The number of tracking points we use\n",
    "        - device: Optional, if specified we use that torch device.\n",
    "        - silent: If False, will print debug info.\n",
    "\n",
    "    Returns:\n",
    "        - torch.Tensor: BTCHW form (where B comes from videos, THW come from latent_mask, and C comes from out_channels)\n",
    "    \"\"\"\n",
    "    dtype = torch.bfloat16\n",
    "    device = device or rp.select_torch_device(prefer_used=True, silent=silent, reserve=True)\n",
    "\n",
    "    LT, LH, LW = latent_mask.shape\n",
    "    LC = out_channels\n",
    "    N = num_points\n",
    "        \n",
    "    videos = rp.as_torch_videos(videos, device=device, dtype=dtype)\n",
    "    B, VT, _, VH, VW = videos.shape\n",
    "\n",
    "    spatial_scale = rp.assert_equality(VH/LH, VW/LW)\n",
    "\n",
    "    #latent_mask -> mask has same dimensions and duration as video, and is bool\n",
    "    mask = latent_mask\n",
    "    mask = rp.resize_list(mask, VT)\n",
    "    mask = rp.as_numpy_array(mask)\n",
    "    mask = rp.as_grayscale_images(mask)\n",
    "    mask = rp.as_binary_images(mask)\n",
    "    mask = rp.resize_images(mask, size=(VH, VW), interp='nearest')\n",
    "    mask = rp.as_binary_images(mask)\n",
    "    assert rp.is_numpy_array(mask)\n",
    "\n",
    "    # Sample N random spatio-temporal points from the mask\n",
    "    track_points = get_random_mask_positions(mask, N)\n",
    "    # Convert to numpy array since run_cotracker expects numpy\n",
    "    track_points = track_points.cpu().numpy() if isinstance(track_points, torch.Tensor) else track_points\n",
    "    track_colors = torch.randn(N, LC, device=device, dtype=dtype)\n",
    "\n",
    "    dotted_latents = []\n",
    "    for video in videos:\n",
    "        #TODO: Can we get away with simply tracking less frames instead of tracking all frames and squashing time afterwards?\n",
    "        tracks, visibility = rp.run_cotracker(\n",
    "            video, \n",
    "            device=device,\n",
    "            queries=track_points,\n",
    "        )\n",
    "\n",
    "        latent_tracks = rp.resize_list(tracks, LT) // spatial_scale #Scale it down spatiotemporally\n",
    "        latent_tracks = latent_tracks.long()\n",
    "        \n",
    "        dotted_latent = torch.zeros(LT, LC, LH, LW, device=device, dtype=dtype)\n",
    "        \n",
    "        # Use optimized scatter_add helper function\n",
    "        dotted_latent = fast_scatter_add(dotted_latent, latent_tracks, track_colors, LT, N, LW, LH)\n",
    "        \n",
    "        dotted_latents.append(dotted_latent)\n",
    "    dotted_latents = torch.stack(dotted_latents)\n",
    "    \n",
    "    rp.validate_tensor_shapes(\n",
    "        videos        ='torch: B VT 3  VH VW',\n",
    "        dotted_latents='torch: B LT LC LH LW',\n",
    "        dotted_latent ='torch:   LT LC LH LW',\n",
    "        video         ='torch:   VT 3  VH VW',\n",
    "        mask          ='numpy:   VT    VH VW',\n",
    "        tracks        ='torch:  VT N XY',\n",
    "        visibility    ='torch:  VT N',\n",
    "        track_points  ='numpy:     N TXY', \n",
    "        latent_tracks ='torch: LT  N XY',\n",
    "        latent_mask   ='       LT    LH LW',\n",
    "        track_colors  ='torch: N  LC',\n",
    "        XY  = 2,\n",
    "        TXY = 3,\n",
    "        **rp.gather_vars('N LC LT LH LW B VT VH VW'), #Already decided\n",
    "        verbose = not silent and 'bold white random green',\n",
    "    )\n",
    "\n",
    "    #CogVideoX Assumptions\n",
    "    assert spatial_scale==8          , \"Assuming we're using CogVideoX\"\n",
    "    assert (LT, LH, LW)==(13, 60, 90), \"Assuming we're using CogVideoX\"\n",
    "    assert (VT, VH, VW)==(49,480,720), \"Assuming we're using CogVideoX\"\n",
    "\n",
    "    return rp.gather_vars('dotted_latents')\n",
    "\n",
    "def demo_dotted_latents(*video_urls):\n",
    "    video_urls = rp.detuple(video_urls) or [\"https://video-previews.elements.envatousercontent.com/23ce1f71-c55d-4bc3-bfad-bc7bf8d8168a/watermarked_preview/watermarked_preview.mp4\"]\n",
    "    videos = rp.load_video(video_urls, use_cache=True)\n",
    "    videos = rp.resize_videos(videos,size=(480,720))\n",
    "    videos = rp.resize_lists(videos, length=49)\n",
    "\n",
    "    rp.tic()\n",
    "    result = generate_dotted_latents(\n",
    "        [video],\n",
    "        np.ones((13, 60, 90)),\n",
    "        out_channels=3,\n",
    "        silent=False,\n",
    "    )\n",
    "    [dotted_latent_video] = result.dotted_latents\n",
    "    rp.ptoc(\"generate_dotted_latents\")\n",
    "    \n",
    "    rp.display_video(\n",
    "        rp.horizontally_concatenated_videos(\n",
    "            rp.resize_list(video, 13),\n",
    "            rp.as_numpy_images(\n",
    "                rp.torch_resize_images(\n",
    "                    dotted_latent_video, size=8, interp=\"nearest\", copy=True\n",
    "                )\n",
    "                / 5\n",
    "                + 0.5\n",
    "            ),\n",
    "        ),\n",
    "        framerate=10,\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "demo_dotted_latents()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python cogvid",
   "language": "python",
   "name": "cogvid"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
